import os
import builtins
import re

from stack_graph_py import StackGraphHelper
from tree_sitter import Language, Parser, Node
from collections import defaultdict

from random_tools.utils import FileUtils, RepoUtils
from constants import Constants


def class_info_path(base_dir: str, repo_name: str, /, *, ext: str) -> str:
    return os.path.join(base_dir, f"{repo_name}_class_info.{ext}")


def function_info_path(base_dir: str, repo_name: str, /, *, ext: str) -> str:
    return os.path.join(base_dir, f"{repo_name}_function_info.{ext}")


def import_info_path(base_dir: str, repo_name: str, /, *, ext: str) -> str:
    return os.path.join(base_dir, f"{repo_name}_import_info.{ext}")


def call_info_path(base_dir: str, repo_name: str, /, *, ext: str) -> str:
    return os.path.join(base_dir, f"{repo_name}_call_info.{ext}")


class FunctionRepoBuilder:
    NODE_DECODE_FORMAT = "utf8"

    def __init__(self, *,
                 repo_base: str,
                 save_base: str,
                 dataset_type: str = "repoeval"):
        self.repo_base: str = repo_base
        self.save_base: str = save_base
        self.dataset_type: str = dataset_type

        match dataset_type:
            case "repoeval" | "cceval":
                self.lang = "python"
                import tree_sitter_python
                ts_lang = Language(tree_sitter_python.language())
            case "cceval-java":
                self.lang = "java"
                import tree_sitter_java
                ts_lang = Language(tree_sitter_java.language())
            case _:
                raise ValueError

        self.parser = Parser(ts_lang)

        self.in_class_name: str | None = None
        self.import_infos: dict[tuple[str, ...], list[dict]] = defaultdict(list[dict])
        self.class_infos: dict[tuple[str, ...], list[dict]] = defaultdict(list[dict])
        self.function_infos: dict[tuple[str, ...], list[dict]] = defaultdict(list[dict])

    def _reset_state(self):
        self.in_class_name = None
        self.import_infos = defaultdict(list[dict])
        self.class_infos = defaultdict(list[dict])
        self.function_infos = defaultdict(list[dict])

    def class_info_path(self, repo_name: str, /, *, ext: str) -> str:
        return class_info_path(self.save_base, repo_name, ext=ext)

    def function_info_path(self, repo_name: str, /, *, ext: str) -> str:
        return function_info_path(self.save_base, repo_name, ext=ext)

    def import_info_path(self, repo_name: str, /, *, ext: str) -> str:
        return import_info_path(self.save_base, repo_name, ext=ext)

    @staticmethod
    def build_import_info(*,
                          start_line: int,
                          end_line: int,
                          import_text: str,
                          fpath_tuple: tuple[str, ...]) -> dict:
        return {
            "start_line": start_line,
            "end_line": end_line,
            "import_text": import_text,
            "fpath_tuple": fpath_tuple,
        }

    @staticmethod
    def build_class_info(*,
                         class_name: str,
                         start_line: int,
                         end_line: int,
                         fpath_tuple: tuple[str, ...]) -> dict:
        return {
            "class_name": class_name,
            "start_line": start_line,
            "end_line": end_line,
            "fpath_tuple": fpath_tuple,
        }

    @staticmethod
    def build_function_info(*,
                            function_name: str,
                            start_line: int,
                            end_line: int,
                            signature_end_line: int,
                            class_method: str | None,
                            fpath_tuple: tuple[str, ...]) -> dict:
        return {
            "function_name": function_name,
            "start_line": start_line,
            "end_line": end_line,
            "signature_end_line": signature_end_line,
            "class_method": class_method,
            "fpath_tuple": fpath_tuple,
        }

    def build_function_repo(self, repo_name: str, /, *,
                            save_jsonl: bool = False,
                            save_pickle: bool = True,
                            file_ext: str = "py"):
        self._reset_state()
        repo_files = RepoUtils.iterate_repo_files(self.repo_base, repo_name, ext=file_ext)
        for file in repo_files:
            fpath_tuple = RepoUtils.get_file_tuple(self.repo_base, file)
            self.build_for_file(fpath_tuple)
        if save_pickle:
            FileUtils.save_pickle(self.class_infos, self.class_info_path(repo_name, ext="pkl"))
            FileUtils.save_pickle(self.function_infos, self.function_info_path(repo_name, ext="pkl"))
            FileUtils.save_pickle(self.import_infos, self.import_info_path(repo_name, ext="pkl"))
        if save_jsonl:
            class_infos_flattened = [
                info for file_info in self.class_infos.values()
                for info in file_info
            ]
            function_infos_flattened = [
                info for file_info in self.function_infos.values()
                for info in file_info
            ]
            import_infos_flattened = [
                info for file_info in self.import_infos.values()
                for info in file_info
            ]
            FileUtils.save_jsonl(class_infos_flattened, self.class_info_path(repo_name, ext="jsonl"))
            FileUtils.save_jsonl(function_infos_flattened, self.function_info_path(repo_name, ext="jsonl"))
            FileUtils.save_jsonl(import_infos_flattened, self.import_info_path(repo_name, ext="jsonl"))

    def build_for_file(self, fpath_tuple: tuple[str, ...]):
        code_lines = RepoUtils.read_lines_by_ftuple(self.repo_base, fpath_tuple)
        max_line = len(code_lines) - 1
        code_str = "".join(code_lines)
        code_bytes = bytes(code_str, self.NODE_DECODE_FORMAT)
        ts_tree = self.parser.parse(code_bytes)
        ts_root = ts_tree.root_node
        self.extract_function_info_from_node(ts_root, fpath_tuple)
        self.extract_global_blocks(fpath_tuple, max_line)

    def extract_global_blocks(self, fpath_tuple: tuple[str, ...], max_line: int):
        class_line_ranges: list[tuple[int, int]] = [
            (info["start_line"], info["end_line"])
            for info in self.class_infos[fpath_tuple]
        ]
        function_line_ranges: list[tuple[int, int]] = [
            (info["start_line"], info["end_line"])
            for info in self.function_infos[fpath_tuple]
        ]
        import_line_ranges: list[tuple[int, int]] = [
            (info["start_line"], info["end_line"])
            for info in self.import_infos[fpath_tuple]
        ]
        all_ranges = class_line_ranges + function_line_ranges + import_line_ranges
        all_ranges.sort(key=lambda x: x[0])
        merged_ranges: list[tuple[int, int]] = []
        for start_line, end_line in all_ranges:
            if len(merged_ranges) == 0 or merged_ranges[-1][1] < start_line:
                merged_ranges.append((start_line, end_line))
            else:
                merged_ranges[-1] = (merged_ranges[-1][0], max(merged_ranges[-1][-1], end_line))

        uncovered_ranges: list[tuple[int, int]] = []
        current_start = 0
        for start_line, end_line in merged_ranges:
            if current_start < start_line:
                uncovered_ranges.append((current_start, start_line - 1))
            current_start = end_line + 1

        if current_start <= max_line:
            uncovered_ranges.append((current_start, max_line))

        global_count: int = 0
        for start_line, end_line in uncovered_ranges:
            global_lines = RepoUtils.read_lines_by_ftuple(self.repo_base,
                                                          fpath_tuple,
                                                          start=start_line,
                                                          end=end_line,
                                                          contains_end=True)
            non_empty: bool = False
            for line in global_lines:
                if line.strip() != "":
                    non_empty = True
                    break

            if non_empty:
                function_info = self.build_function_info(function_name=f"$global_{global_count}",
                                                         start_line=start_line,
                                                         end_line=end_line,
                                                         signature_end_line=end_line,
                                                         class_method=None,
                                                         fpath_tuple=fpath_tuple)
                self.function_infos[fpath_tuple].append(function_info)
                global_count += 1

    def extract_function_info_from_node(self, node: Node, fpath_tuple: tuple[str, ...]):
        in_class_flag: bool = False
        match node.type:
            case "import_from_statement" | "import_statement":
                start_line = node.start_point[0]
                end_line = node.end_point[0]
                import_text = node.text.decode(self.NODE_DECODE_FORMAT)
                import_info = self.build_import_info(start_line=start_line,
                                                     end_line=end_line,
                                                     import_text=import_text,
                                                     fpath_tuple=fpath_tuple)
                self.import_infos[fpath_tuple].append(import_info)
            case "class_definition":
                if self.in_class_name is not None:
                    # ignore nested class
                    return
                start_row = node.start_point[0]
                end_row = node.end_point[0]
                class_name = node.child_by_field_name("name").text.decode(self.NODE_DECODE_FORMAT)
                class_info = self.build_class_info(class_name=class_name,
                                                   start_line=start_row,
                                                   end_line=end_row,
                                                   fpath_tuple=fpath_tuple)

                self.class_infos[fpath_tuple].append(class_info)
                self.in_class_name = class_name
                in_class_flag = True
            case "function_definition":
                start_row = node.start_point[0]
                end_row = node.end_point[0]
                signature_end_row = node.child_by_field_name("parameters").end_point[0]

                function_name = node.child_by_field_name("name").text.decode(self.NODE_DECODE_FORMAT)
                class_method = self.in_class_name

                function_info = self.build_function_info(function_name=function_name,
                                                         start_line=start_row,
                                                         end_line=end_row,
                                                         signature_end_line=signature_end_row,
                                                         class_method=class_method,
                                                         fpath_tuple=fpath_tuple)

                self.function_infos[fpath_tuple].append(function_info)

                # ignore nested function
                return
            case _:
                pass

        for child in node.children:
            self.extract_function_info_from_node(child, fpath_tuple)

        if in_class_flag:
            # traverse class finished, remove in_class_name
            self.in_class_name = None
